# coding: utf-8
from netCDF4 import Dataset
import matplotlib.pyplot as plt
import Gauge


def Deal():
    centers = []
    centerX0 = GetCenter(
        r'F:/Research/WRF_Output/2008/01/2018/60/wrfout_d03_2008-01-20_18_00_00')
    centerX1 = GetCenter(
        r'F:/Research/WRF_Output/2008/01/2018/61/wrfout_d03_2008-01-20_18_00_00')
    centerX2 = GetCenter(
        r'F:/Research/WRF_Output/2008/01/2018/62/wrfout_d03_2008-01-20_18_00_00')
    centerX10 = GetCenter(
        r'F:/Research/WRF_Output/2008/01/2018/610/wrfout_d03_2008-01-20_18_00_00')

    centerX0 = decumulation(centerX0)
    centerX1 = decumulation(centerX1)
    centerX2 = decumulation(centerX2)
    centerX10 = decumulation(centerX10)

    centers.append(centerX0)
    centers.append(centerX1)
    centers.append(centerX2)
    centers.append(centerX10)

    mean = Gauge.Get50GaugeMean("2008-01-20_18:00:00", "2008-01-21_12:00:00")
    centers.append(mean)

    labels = ('20', '21', '22', '210', 'gague')
    # labels = ('20', '21', '22', '210')

    PlotCenters(centers, labels)


def GetCenter(source):
    '''deal with the output of WRF'''

    if source:
        rainfall = Dataset(source)
        rainc = rainfall.variables['RAINC']
        rainnc = rainfall.variables['RAINNC']
        (Time_Number, SN_Number, WE_Number) = rainnc.shape
        centerX = SN_Number // 2
        centerY = WE_Number // 2
        center = []
        for t in range(0, Time_Number):
            # get the rainfall data according to time
            rainc_2D = rainc[t, :, :]
            rainnc_2D = rainnc[t, :, :]
            precipitation = rainc_2D[centerX, centerY] + rainnc_2D[centerX, centerY]
            center.append(precipitation)
        rainfall.close()
        return center


def PlotCenters(centers, labels):
    '''plot rainfall data of center'''
    time = range(len(centers[0]))
    plt.figure(figsize=(10, 10))
    for i in range(len(centers)):
        plt.plot(time, centers[i], label=labels[i])

    plt.legend(loc='best', fontsize='xx-large')
    plt.show()


def decumulation(seq):
    newSeq = []
    newSeq.append(seq[0])
    for i in range(len(seq) - 1):
        newSeq.append(seq[i + 1] - seq[i])
    return newSeq


if __name__ == '__main__':
    Deal()
